from datasets.basedataset import BaseDataset
import torch.utils.data as data
from nets.bases.basicmodel import BasicModule
import os
import shutil

import torch as t
import cv2 as cv
import numpy as np
from torch import autograd
from torch.functional import Tensor
import torch.utils.data as data
from tqdm import tqdm
from torch.autograd import Variable
from torch.utils.tensorboard import SummaryWriter
import torchvision
import shutil
from data_enhancement.resize_convert import ResizeConvert
from torchstat import stat

class TrainBase(object):
    
    def create_train(self, dataset : BaseDataset, model : BasicModule, batch_size, shuffle = True, num_workers = 2, size = (512, 512), optim = 'Adam', lr = 1e-3,
                 save_model_count = 100, device = 'cuda:0', is_print_model_info = False):
        #模型
        self.model : BasicModule = model
        print('is_print_model_info : ', is_print_model_info)
        if is_print_model_info == True:
            print('model_info : \n')
            stat(self.model, self.get_model_test_size())
            raise Exception('模型信息打印完毕，正常运行请关闭打印信息功能!')
         #训练数据集
        self.dataset = data.DataLoader(dataset, batch_size = batch_size, shuffle = shuffle, num_workers = num_workers)
        #日志路径
        self.log_path = './logs/'
        #保存模型路径
        self.save_model_path = './models/'
        #保存模型的迭代次数
        self.save_model_count = save_model_count;
        #优化器
        print('optim : ', optim)
        if optim == 'Adam':
            self.optim = t.optim.Adam(self.model.parameters(), lr=lr)
        elif optim == 'SGD':
            self.optim = t.optim.SGD(self.model.parameters(), lr=lr)
        elif optim == 'RMS':
            self.optim = t.optim.RMSprop(self.model.parameters(), lr=lr)
            
        self.device = t.device(device)
        #转换到对应的设备平台运行
        self.model.to(self.device)
        
        print('params : \n', 
              'learning rate : ', lr, '\n',
              'batch size : ', batch_size, '\n',
              'shuffle : ', shuffle, '\n',
              'num_workers : ', num_workers, '\n',
              'size : ', size, '\n',
              'optim : ', optim, '\n',
              'device : ', self.device.type,
              '\n')
        
        #resize
        self.resize_convert = ResizeConvert(size)
    
    #创建测试用模块
    def create_test(self, model : BasicModule, size = (512, 512), device='cuda:0', model_path = None):
        #resize
        self.resize_convert = ResizeConvert(size)
        self.device = t.device(device)
        #模型
        self.model : BasicModule = model
        self.model.loadAny(model_path)
        self.model.to(self.device)
        self.model.eval()
        
    
    #获取模型测试的图片size  不包含batchsize
    def get_model_test_size(self):
        pass
        
    
    #单次训练
    def train_once(self, datas):
        pass
    
    #计算损失
    def caculator_loss(self, res):
        pass
    
    #训练
    def train(self, epochs = 100):
        self.model.train()
        data_len = self.dataset.__len__();
        print('开始训练……………… 数据长度 : ', data_len)
        # 删除日志文件，重新创建
        if os.path.exists(self.log_path) == True:
            shutil.rmtree(self.log_path)
            os.mkdir(self.log_path)
        
        #tensorboard 日志
        write = SummaryWriter(self.log_path)

        #训练
        for epoch in range(epochs):
            loss_total = []
            pbar = tqdm(
                total=data_len, desc='epoch : ' + str(epoch))
            for step, data in enumerate(self.dataset):
                res = self.train_once(data)
                
                #清空梯度，防止梯度累积，反向传播
                self.model.zero_grad()
                self.optim.zero_grad()
                loss = self.caculator_loss(res)
                loss.backward()
                self.optim.step()
                
                loss_total.append(loss.item())
                pbar.update(1)
            
            pbar.close()
            self.model.save(epoch)
            loss_mean = np.mean(loss_total)
            write.add_scalar('loss', loss_mean, epoch)
            print('loss : ', loss_mean, ' epoch : ', epoch)
            

    #预测
    def test(self, path, visualize = False):
        img = self.load_img(path)
        if visualize == True:
            res = self.model.visualize(img)
        else:
            res = self.model(img)
        res = t.sigmoid(res)
        res = res * 255.0
        res = t.squeeze(res)
        res[res < 128] = 0
        return res.detach().numpy() if self.device.type == 'cpu' else res.detach().cpu().numpy()
    
    #加载模型
    def load(self, path):
        print('load model : ', path)
        self.model.loadAny(path)
        
        
    #加载图片
    def load_img(self, path):
        img = cv.imdecode(np.fromfile(path, dtype=np.uint8), -1)
        img,_, _ = self.resize_convert.convert(img, None, None)
        img = img.transpose((2, 0, 1)).astype(np.float32) / 255.0
        img = img[np.newaxis, :, :, :]
        img = t.from_numpy(img).to(self.device)
        return img